import asyncio
import websockets
import json
import uuid
from datetime import datetime
import time
import requests
import logging

# 配置测试日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class WebSocketServerTester:
    """WebSocket 服务器测试类"""

    def __init__(self, base_url="http://localhost:8000", ws_url="ws://localhost:8000"):
        self.base_url = base_url
        self.ws_url = ws_url
        self.connected_devices = {}

    async def test_device_registration(self, device_count=3):
        """测试设备注册功能"""
        logger.info(f"🧪 测试设备注册 - 准备注册 {device_count} 个设备")

        tasks = []
        for i in range(device_count):
            device_id = f"test_device_{i + 1}"
            task = asyncio.create_task(self._register_single_device(device_id))
            tasks.append(task)

        results = await asyncio.gather(*tasks, return_exceptions=True)

        success_count = sum(1 for result in results if result)
        logger.info(f"✅ 设备注册测试完成: {success_count}/{device_count} 成功")
        return success_count

    async def _register_single_device(self, device_id: str):
        """注册单个设备"""
        try:
            async with websockets.connect(f"{self.ws_url}/ws/device/{device_id}") as websocket:
                # 发送注册消息
                register_message = {
                    "type": "device_register",
                    "data": {
                        "deviceId": device_id,
                        "versionName": "v1.0.0",
                        "url": f"ws://client/{device_id}",
                        "code": f"CODE_{device_id.upper()}",
                        "capabilities": ["temperature", "humidity", "pressure"]
                    },
                    "timestamp": datetime.now().isoformat()
                }

                await websocket.send(json.dumps(register_message))

                # 等待响应
                response = await asyncio.wait_for(websocket.recv(), timeout=10)
                response_data = json.loads(response)

                if response_data.get("type") == "status_update":
                    self.connected_devices[device_id] = websocket
                    logger.info(f"✅ 设备 {device_id} 注册成功")
                    return True
                else:
                    logger.error(f"❌ 设备 {device_id} 注册失败: {response_data}")
                    return False

        except Exception as e:
            logger.error(f"❌ 设备 {device_id} 连接失败: {e}")
            return False

    async def test_heartbeat(self, device_id: str = "test_device_1"):
        """测试心跳功能"""
        logger.info(f"🧪 测试心跳功能 - 设备: {device_id}")

        try:
            if device_id not in self.connected_devices:
                await self._register_single_device(device_id)

            websocket = self.connected_devices[device_id]

            # 发送心跳消息
            heartbeat_message = {
                "type": "heartbeat",
                "data": {
                    "status": "healthy",
                    "cpu_usage": 45.2,
                    "memory_usage": 67.8,
                    "uptime": 3600
                },
                "timestamp": datetime.now().isoformat()
            }

            await websocket.send(json.dumps(heartbeat_message))

            # 等待心跳响应
            response = await asyncio.wait_for(websocket.recv(), timeout=10)
            response_data = json.loads(response)

            if response_data.get("type") == "heartbeat":
                logger.info(f"✅ 设备 {device_id} 心跳测试成功")
                return True
            else:
                logger.error(f"❌ 设备 {device_id} 心跳测试失败: {response_data}")
                return False

        except Exception as e:
            logger.error(f"❌ 设备 {device_id} 心跳测试异常: {e}")
            return False

    async def test_status_update(self, device_id: str = "test_device_1"):
        """测试状态更新功能"""
        logger.info(f"🧪 测试状态更新 - 设备: {device_id}")

        try:
            if device_id not in self.connected_devices:
                await self._register_single_device(device_id)

            websocket = self.connected_devices[device_id]

            # 发送状态更新消息
            status_message = {
                "type": "status_update",
                "data": {
                    "temperature": 25.5,
                    "humidity": 65.2,
                    "pressure": 1013.25,
                    "battery_level": 85,
                    "network_strength": 4
                },
                "timestamp": datetime.now().isoformat()
            }

            await websocket.send(json.dumps(status_message))

            # 状态更新不需要响应，等待一下确保消息发送完成
            await asyncio.sleep(1)
            logger.info(f"✅ 设备 {device_id} 状态更新测试成功")
            return True

        except Exception as e:
            logger.error(f"❌ 设备 {device_id} 状态更新测试异常: {e}")
            return False

    async def test_data_report(self, device_id: str = "test_device_1"):
        """测试数据报告功能"""
        logger.info(f"🧪 测试数据报告 - 设备: {device_id}")

        try:
            if device_id not in self.connected_devices:
                await self._register_single_device(device_id)

            websocket = self.connected_devices[device_id]

            # 发送数据报告消息
            data_message = {
                "type": "data_report",
                "data": {
                    "data_points": [
                        {"timestamp": datetime.now().isoformat(), "value": 25.1, "sensor": "temp"},
                        {"timestamp": datetime.now().isoformat(), "value": 65.3, "sensor": "humidity"},
                        {"timestamp": datetime.now().isoformat(), "value": 1013.2, "sensor": "pressure"}
                    ],
                    "total_count": 3,
                    "sampling_rate": "1s"
                },
                "timestamp": datetime.now().isoformat()
            }

            await websocket.send(json.dumps(data_message))
            await asyncio.sleep(1)
            logger.info(f"✅ 设备 {device_id} 数据报告测试成功")
            return True

        except Exception as e:
            logger.error(f"❌ 设备 {device_id} 数据报告测试异常: {e}")
            return False

    async def test_command_receiving(self, device_id: str = "test_device_1"):
        """测试接收服务器命令"""
        logger.info(f"🧪 测试命令接收 - 设备: {device_id}")

        try:
            if device_id not in self.connected_devices:
                await self._register_single_device(device_id)

            websocket = self.connected_devices[device_id]

            # 设置一个消息接收任务
            async def receive_command():
                try:
                    message = await asyncio.wait_for(websocket.recv(), timeout=35)
                    message_data = json.loads(message)
                    if message_data.get("type") == "command":
                        logger.info(f"✅ 设备 {device_id} 成功接收到命令: {message_data}")
                        return True
                except asyncio.TimeoutError:
                    logger.warning(f"⚠️  设备 {device_id} 未在预期时间内收到命令")
                    return False

            # 等待接收命令（服务器应该每30秒发送一次状态查询）
            result = await receive_command()
            return result

        except Exception as e:
            logger.error(f"❌ 设备 {device_id} 命令接收测试异常: {e}")
            return False

    def test_rest_api(self):
        """测试 REST API 接口"""
        logger.info("🧪 测试 REST API 接口")

        try:
            # 测试根路径
            response = requests.get(f"{self.base_url}/")
            if response.status_code == 200:
                logger.info("✅ 根路径测试成功")
            else:
                logger.error(f"❌ 根路径测试失败: {response.status_code}")

            # 测试获取在线设备
            response = requests.get(f"{self.base_url}/devices/online")
            if response.status_code == 200:
                online_devices = response.json()
                logger.info(f"✅ 在线设备查询成功: {len(online_devices)} 个设备在线")
            else:
                logger.error(f"❌ 在线设备查询失败: {response.status_code}")

            # 测试获取特定设备状态
            response = requests.get(f"{self.base_url}/devices/test_device_1")
            if response.status_code == 200:
                device_status = response.json()
                logger.info(f"✅ 设备状态查询成功: {device_status}")
            elif response.status_code == 404:
                logger.warning("⚠️  设备状态查询 - 设备不存在（正常情况）")
            else:
                logger.error(f"❌ 设备状态查询失败: {response.status_code}")

            return True

        except Exception as e:
            logger.error(f"❌ REST API 测试异常: {e}")
            return False

    async def test_broadcast_functionality(self):
        """测试广播功能"""
        logger.info("🧪 测试广播功能")

        try:
            # 首先注册几个设备
            await self.test_device_registration(2)

            # 通过 REST API 发送广播消息
            broadcast_data = {
                "message": "系统维护通知",
                "type": "maintenance",
                "start_time": datetime.now().isoformat(),
                "duration": "30分钟"
            }

            response = requests.post(
                f"{self.base_url}/broadcast",
                json=broadcast_data
            )

            if response.status_code == 200:
                logger.info("✅ 广播消息发送成功")
                return True
            else:
                logger.error(f"❌ 广播消息发送失败: {response.status_code}")
                return False

        except Exception as e:
            logger.error(f"❌ 广播功能测试异常: {e}")
            return False

    async def test_error_handling(self):
        """测试错误处理"""
        logger.info("🧪 测试错误处理")

        try:
            # 测试发送无效消息
            async with websockets.connect(f"{self.ws_url}/ws/device/error_test") as websocket:
                invalid_message = {
                    "type": "invalid_type",  # 无效的消息类型
                    "data": {"test": "data"}
                }

                await websocket.send(json.dumps(invalid_message))

                # 应该收到错误响应
                response = await asyncio.wait_for(websocket.recv(), timeout=10)
                response_data = json.loads(response)

                if response_data.get("type") == "error":
                    logger.info("✅ 错误处理测试成功 - 正确返回错误响应")
                    return True
                else:
                    logger.error(f"❌ 错误处理测试失败: {response_data}")
                    return False

        except Exception as e:
            logger.error(f"❌ 错误处理测试异常: {e}")
            return False

    async def run_comprehensive_test(self):
        """运行全面测试"""
        logger.info("🚀 开始全面测试 WebSocket 服务器")

        test_results = {}

        # 1. 测试设备注册
        test_results["device_registration"] = await self.test_device_registration(2)

        # 2. 测试心跳功能
        test_results["heartbeat"] = await self.test_heartbeat()

        # 3. 测试状态更新
        test_results["status_update"] = await self.test_status_update()

        # 4. 测试数据报告
        test_results["data_report"] = await self.test_data_report()

        # 5. 测试命令接收
        test_results["command_receiving"] = await self.test_command_receiving()

        # 6. 测试 REST API
        test_results["rest_api"] = self.test_rest_api()

        # 7. 测试广播功能
        test_results["broadcast"] = await self.test_broadcast_functionality()

        # 8. 测试错误处理
        test_results["error_handling"] = await self.test_error_handling()

        # 汇总结果
        logger.info("\n" + "=" * 50)
        logger.info("📊 测试结果汇总")
        logger.info("=" * 50)

        for test_name, result in test_results.items():
            status = "✅ 通过" if result else "❌ 失败"
            logger.info(f"{test_name:20} {status}")

        total_tests = len(test_results)
        passed_tests = sum(1 for result in test_results.values() if result)

        logger.info("=" * 50)
        logger.info(f"总测试数: {total_tests}, 通过: {passed_tests}, 失败: {total_tests - passed_tests}")

        if passed_tests == total_tests:
            logger.info("🎉 所有测试通过！")
        else:
            logger.warning("⚠️  部分测试失败，请检查服务器配置")

        return test_results

    async def cleanup(self):
        """清理测试资源"""
        logger.info("🧹 清理测试资源")
        for device_id, websocket in self.connected_devices.items():
            await websocket.close()
        self.connected_devices.clear()


# 使用示例
async def main():
    """主测试函数"""
    tester = WebSocketServerTester()

    try:
        # 运行全面测试
        await tester.run_comprehensive_test()

        # 或者运行单个测试
        # await tester.test_device_registration(3)
        # await tester.test_heartbeat()
        # tester.test_rest_api()

    except Exception as e:
        logger.error(f"测试过程中发生错误: {e}")
    finally:
        await tester.cleanup()


if __name__ == "__main__":
    # 运行测试
    asyncio.run(main())